#include "G4SystemOfUnits.hh"
#include "Randomize.hh"
#include "G4ParticleTable.hh"
#include "G4IonTable.hh"
#include "G4ParticleDefinition.hh"
#include "G4HadronElastic.hh"
#include "G4Proton.hh"
#include "G4Neutron.hh"
#include "G4Deuteron.hh"
#include "G4Alpha.hh"
#include "G4Triton.hh"
#include "G4He3.hh"
#include "G4Pow.hh"
#include <memory>

#include "MyGenCreator.hh"
#include "MyGenContainer.hh"

#include "NuclearRecoil.hh"

NuclearRecoil::NuclearRecoil(): MyBaseGen()
{
    m_PGType = "Nuclear Recoil Background";
    m_ClassName = "NuclearRecoil";
    m_TypeFlag = 1;
    m_EvtFlag = 0;

    m_Particle = "C+";
    m_ParticleEnergy = 1. * MeV;
    m_ParticlePosition = G4ThreeVector(0., 0., 0.);
    m_ParticlePolarization = G4ThreeVector(0., 0., 0.);
    m_ParticleMomentumDirection = G4ThreeVector(0., 0., 0.);

    m_Gun = nullptr;
    m_PosGen = nullptr;
    m_IonPD = nullptr;
    m_NeutronPD = nullptr;
}

NuclearRecoil::~NuclearRecoil()
{
    if(!m_Gun)
        delete m_Gun;
    if(!m_FileEkin)
        delete m_FileEkin;
}

void NuclearRecoil::Initialize(YAML::Node yamlNode)
{
    m_Gun = new G4ParticleGun(1);
    m_PosGen = MyGenContainer::GetInstance().GetMyPos();

    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    m_NeutronPD = particleTable->FindParticle("neutron");

    m_Z = yamlNode["Z"].as<int>(6); // 原子序数
    m_A = yamlNode["A"].as<double>(12); // 质量数
    m_Charge = yamlNode["Charge"].as<double>(-1e6); // 电荷
    if(m_Charge <= -1e6)
    {
        m_Charge = m_Z;
    }

    m_FileEkin = new TFile(yamlNode["HistFile"].as<std::string>("NeutronKineticEnergy.root").c_str());
    m_HistEkin = dynamic_cast<TH1D*>(m_FileEkin->Get(yamlNode["HistName"].as<std::string>("NeutronEkin").c_str()));
}

void NuclearRecoil::GeneratePrimaryVertex(G4Event* anEvent)
{
    if(m_IonPD == nullptr)
    {
        if(m_Z == 1 && m_A == 1){
            m_IonPD   = G4Proton::Proton();
            m_Particle = m_IonPD->GetParticleName();
        }
        else if(m_Z == 1 && m_A == 2)
        {
            m_IonPD = G4Deuteron::Deuteron();
            m_Particle = m_IonPD->GetParticleName();
        }
        else if(m_Z == 1 && m_A == 3)
        {
            m_IonPD = G4Triton::Triton();
            m_Particle = m_IonPD->GetParticleName();
        }
        else if(m_Z == 2 && m_A == 3)
        {
            m_IonPD = G4He3::He3();
            m_Particle = m_IonPD->GetParticleName();
        }
        else if(m_Z == 2 && m_A == 4)
        {
            m_IonPD = G4Alpha::Alpha();
            m_Particle = m_IonPD->GetParticleName();
        }
        else{
            G4IonTable *ionTable = (G4IonTable *)(G4ParticleTable::GetParticleTable()->GetIonTable());
            m_IonPD = ionTable->GetIon(m_Z, m_A);
            m_Particle = ionTable->GetIonName(m_Z, m_A);
        }
        m_Gun->SetParticleDefinition(m_IonPD);
        m_Gun->SetParticleCharge(m_Charge);
        std::cout << m_Particle << "\t Charge: " << m_Charge << std::endl;
    }

    PosAndMom neutronPosAndMomDir = m_PosGen->GenPosAndMomDir();
    G4double neutronEk = m_HistEkin->GetRandom() * MeV;
    G4DynamicParticle neutronTrack(m_NeutronPD, neutronPosAndMomDir.second, neutronEk);
    G4DynamicParticle ionTrack = NuclearRecoil::HadronNuclearScattering(neutronTrack, m_Z, m_A);
    m_ParticleEnergy = ionTrack.GetKineticEnergy();
    m_ParticlePosition = neutronPosAndMomDir.first;
    m_ParticleMomentumDirection = ionTrack.GetMomentumDirection();

    m_Gun->SetParticleEnergy(m_ParticleEnergy);
    m_Gun->SetParticlePosition(neutronPosAndMomDir.first);
    m_Gun->SetParticleMomentumDirection(m_ParticleMomentumDirection);
    m_Gun->GeneratePrimaryVertex(anEvent);
    ++m_EvtFlag;
}

G4DynamicParticle NuclearRecoil::HadronNuclearScattering(G4DynamicParticle& hadron, G4int Z, G4int A)
// Refer to G4HadronElastic::ApplyYourself()
{
    // Scattered particle referred to axis of incident particle
    G4double kineticEnergy = hadron.GetKineticEnergy();
    G4double mass = hadron.GetMass();
    G4double plab = std::sqrt(kineticEnergy * (kineticEnergy + 2.0 * mass));

    G4double mass2 = G4NucleiProperties::GetNuclearMass(A, Z);
    G4double e1 = mass + kineticEnergy;
    // Assuming that the hadron injected from positice z-axis
    G4LorentzVector lv(0.0, 0.0, plab, e1 + mass2);
    G4ThreeVector bst = lv.boostVector();
    G4double momentumCMS = plab * mass2 / std::sqrt(mass * mass + mass2 * mass2 + 2. * mass2 * e1);

    G4double localTmax = 4.0 * momentumCMS * momentumCMS;

    // Sampling in CM system
    G4double t = NuclearRecoil::SampleInvariantT(localTmax, A);

    G4double phi = G4UniformRand() * CLHEP::twopi;
    G4double cost = 1. - 2.0 * t / localTmax;

    if (cost > 1.0)
    {
        cost = 1.0;
    }
    else if (cost < -1.0)
    {
        cost = -1.0;
    }

    G4double sint = std::sqrt((1.0 - cost) * (1.0 + cost));

    G4LorentzVector nlv1(momentumCMS * sint * std::cos(phi),
                         momentumCMS * sint * std::sin(phi),
                         momentumCMS * cost,
                         std::sqrt(momentumCMS * momentumCMS + mass * mass));

    nlv1.boost(bst);

    lv -= nlv1;
    G4double erec = std::max(lv.e() - mass2, 0.0);

    // Rotate to LAB reference
    G4ThreeVector hadronMomDir = hadron.GetMomentumDirection();
    G4ThreeVector rotateVect = G4ThreeVector(0., 0., 1.).cross(hadronMomDir);
    G4double rotateAngle = G4ThreeVector(0., 0., 1.).angle(hadronMomDir);
    lv.rotate(rotateAngle, rotateVect);

    // the recoil is created if kinetic energy above the threshold
    G4DynamicParticle ion(G4ParticleTable::GetParticleTable()->GetIonTable()->GetIon(Z,A,0.0), lv.vect().unit(), erec);

    return ion;
}

G4double NuclearRecoil::SampleInvariantT(G4double localTmax, G4int A)
// Refer to G4HadronElastic::SampleInvariantT()
{
    const G4double GeV2 = GeV * GeV;
    const G4double numLimit = 18.;

    G4double tmax = localTmax / GeV2;

    G4double aa, bb, cc, dd;
    G4Pow *g4pow = G4Pow::GetInstance();
    if (A <= 62)
    {
        bb = 14.5 * g4pow->Z23(A);
        dd = 20.;
        aa = (A * A) / bb; // 1.63
        cc = 1.4 * g4pow->Z13(A) / dd;
    }
    else
    { //(A>62)
        bb = 60. * g4pow->Z13(A);
        dd = 25.;
        aa = g4pow->powZ(A, 1.33) / bb;      // 1.33
        cc = 0.2 * g4pow->powZ(A, 0.4) / dd; // 1:0.4     ---    2: 0.4
    }
    G4double q1 = 1.0 - G4Exp(-std::min(bb * tmax, numLimit));
    G4double q2 = 1.0 - G4Exp(-std::min(dd * tmax, numLimit));
    G4double s1 = q1 * aa;
    G4double s2 = q2 * cc;
    if ((s1 + s2) * G4UniformRand() < s2)
    {
        q1 = q2;
        bb = dd;
    }

    G4double invT = -1.;
    while(invT < 0.0 || invT > localTmax)
    {
        invT = -GeV2 * G4Log(1.0 - G4UniformRand() * q1) / bb;
    }
    return invT;
}

REGISTGPSGEN(NuclearRecoil)